use std::{
    fmt,
    sync::{Arc, Mutex},
};

use baml_types::{BamlValueWithMeta, Completion, Constraint, ResponseCheck, TypeIR};

use crate::hir::HeaderContext;

/// Unique identifier for a streaming watch notification
pub type StreamId = String;

/// Concrete handler for watch notifications.
/// This replaces the trait-based approach with conditional Send bounds.
/// Use Arc<Mutex<WatchHandler>> for shared ownership across threads.
pub struct WatchHandler<F>
where
    F: FnMut(WatchNotification),
{
    handler: F,
}

impl<F> WatchHandler<F>
where
    F: FnMut(WatchNotification),
{
    pub fn new(handler: F) -> Self {
        WatchHandler { handler }
    }

    pub fn notify(&mut self, notification: WatchNotification) {
        (self.handler)(notification)
    }
}

#[cfg(not(target_arch = "wasm32"))]
impl WatchHandler<Box<dyn FnMut(WatchNotification) + Send>> {
    /// Create a no-op handler (native)
    pub fn noop() -> Self {
        WatchHandler {
            handler: Box::new(|_| {}),
        }
    }
}

#[cfg(target_arch = "wasm32")]
impl WatchHandler<Box<dyn FnMut(WatchNotification)>> {
    /// Create a no-op handler (WASM)
    pub fn noop() -> Self {
        WatchHandler {
            handler: Box::new(|_| {}),
        }
    }
}

/// Type alias for shared watch handler (thread-safe for native, single-threaded for WASM)
#[cfg(not(target_arch = "wasm32"))]
pub type SharedWatchHandler = Arc<Mutex<WatchHandler<Box<dyn FnMut(WatchNotification) + Send>>>>;

#[cfg(target_arch = "wasm32")]
pub type SharedWatchHandler = Arc<Mutex<WatchHandler<Box<dyn FnMut(WatchNotification)>>>>;

/// Helper to create a shared watch handler from a function (native - requires Send)
#[cfg(not(target_arch = "wasm32"))]
pub fn shared_handler<F>(handler: F) -> SharedWatchHandler
where
    F: FnMut(WatchNotification) + Send + 'static,
{
    Arc::new(Mutex::new(WatchHandler::new(Box::new(handler))))
}

/// Helper to create a shared watch handler from a function (WASM - no Send requirement)
#[cfg(target_arch = "wasm32")]
pub fn shared_handler<F>(handler: F) -> SharedWatchHandler
where
    F: FnMut(WatchNotification) + 'static,
{
    #[allow(clippy::arc_with_non_send_sync)]
    Arc::new(Mutex::new(WatchHandler::new(Box::new(handler))))
}

/// Helper to create a shared no-op handler
pub fn shared_noop_handler() -> SharedWatchHandler {
    #[allow(clippy::arc_with_non_send_sync)]
    Arc::new(Mutex::new(WatchHandler::noop()))
}

#[derive(Debug, Clone)]
pub enum WatchBamlValue {
    Value(BamlValueWithMeta<WatchValueMetadata>),
    Header(HeaderContext),
    /// HACK: Emitted synthetically when a new header comes in at the same or shallower level.
    /// This signals that the previous header's scope has ended. This is a workaround until
    /// proper exit events are emitted from the interpreter.
    HeaderStopped(HeaderContext),
    StreamStart(StreamId),
    StreamUpdate(StreamId, BamlValueWithMeta<WatchValueMetadata>),
    StreamEnd(StreamId),
}

/// The BamlValueWithMeta metadata for a
/// BamlValue in a notification.
#[derive(Debug, Clone)]
pub struct WatchValueMetadata {
    pub constraints: Vec<Constraint>,
    pub response_checks: Vec<ResponseCheck>,
    pub completion: Completion,
    pub r#type: TypeIR,
}

#[derive(Debug)]
pub struct WatchNotification {
    pub value: WatchBamlValue,
    pub variable_name: Option<String>,
    pub channel_name: Option<String>,
    pub function_name: String,
    pub is_stream: bool,
}

impl fmt::Display for WatchNotification {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self.value {
            WatchBamlValue::Value(value) => match (&self.variable_name, &self.channel_name) {
                (Some(var_name), Some(chan_name)) if var_name != chan_name => {
                    write!(
                        f,
                        "(var) {} [channel: {}]: {}",
                        var_name,
                        chan_name,
                        value.clone().value()
                    )
                }
                (Some(var_name), _) => {
                    write!(f, "(var) {}: {}", var_name, value.clone().value())
                }
                (None, Some(chan_name)) => {
                    write!(f, "(channel) {}: {}", chan_name, value.clone().value())
                }
                _ => {
                    write!(f, "{}", value.clone().value())
                }
            },
            WatchBamlValue::Header(header) => {
                write!(
                    f,
                    "(header L{level}) {function}.{title}",
                    level = header.level,
                    function = self.function_name,
                    title = header.title
                )
            }
            WatchBamlValue::HeaderStopped(header) => {
                write!(
                    f,
                    "(header stopped L{level}) {function}.{title}",
                    level = header.level,
                    function = self.function_name,
                    title = header.title
                )
            }
            WatchBamlValue::StreamStart(stream_id) => {
                write!(f, "(stream start) {stream_id}")
            }
            WatchBamlValue::StreamUpdate(stream_id, value) => {
                write!(
                    f,
                    "(stream update) {}: {}",
                    stream_id,
                    value.clone().value()
                )
            }
            WatchBamlValue::StreamEnd(stream_id) => {
                write!(f, "(stream end) {stream_id}")
            }
        }
    }
}

impl WatchNotification {
    pub fn new_var(
        variable_name: String,
        channel_name: String,
        value: BamlValueWithMeta<WatchValueMetadata>,
        function_name: String,
    ) -> Self {
        Self {
            value: WatchBamlValue::Value(value),
            variable_name: Some(variable_name),
            channel_name: Some(channel_name),
            function_name,
            is_stream: false,
        }
    }

    pub fn new_stream(
        variable_name: String,
        value: BamlValueWithMeta<WatchValueMetadata>,
        function_name: String,
    ) -> Self {
        Self {
            value: WatchBamlValue::Value(value),
            variable_name: Some(variable_name.clone()),
            channel_name: Some(variable_name),
            function_name,
            is_stream: true,
        }
    }

    pub fn new_block(header: HeaderContext, function_name: String) -> Self {
        Self {
            value: WatchBamlValue::Header(header),
            variable_name: None,
            channel_name: None,
            function_name,
            is_stream: false,
        }
    }

    /// HACK: Create a synthetic "stopped" notification for a header block.
    /// Used when a new header at the same or shallower level comes in.
    pub fn new_block_stopped(header: HeaderContext, function_name: String) -> Self {
        Self {
            value: WatchBamlValue::HeaderStopped(header),
            variable_name: None,
            channel_name: None,
            function_name,
            is_stream: false,
        }
    }

    pub fn new_stream_start(
        variable_name: String,
        stream_id: StreamId,
        function_name: String,
    ) -> Self {
        Self {
            value: WatchBamlValue::StreamStart(stream_id),
            variable_name: Some(variable_name.clone()),
            channel_name: Some(variable_name),
            function_name,
            is_stream: true,
        }
    }

    pub fn new_stream_update(
        variable_name: String,
        stream_id: StreamId,
        value: BamlValueWithMeta<WatchValueMetadata>,
        function_name: String,
    ) -> Self {
        Self {
            value: WatchBamlValue::StreamUpdate(stream_id, value),
            variable_name: Some(variable_name.clone()),
            channel_name: Some(variable_name),
            function_name,
            is_stream: true,
        }
    }

    pub fn new_stream_end(
        variable_name: String,
        stream_id: StreamId,
        function_name: String,
    ) -> Self {
        Self {
            value: WatchBamlValue::StreamEnd(stream_id),
            variable_name: Some(variable_name.clone()),
            channel_name: Some(variable_name),
            function_name,
            is_stream: true,
        }
    }
}
